package bsh;

import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Map.Entry;

import jmine.tec.utils.Tuple;
import jmine.tec.utils.UnmodifiableTuple;

import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;

import bsh.exception.CompilationAssertionFailed;
import bsh.exception.CompilationFailedException;
import bsh.meta.JavaClassWrapper;
import bsh.meta.Scope;
import bsh.meta.ScopeType;
import bsh.meta.TypeWrapper;

/**
 * Compiles a try { ... } catch (....) { .... } finally { .... } statement
 * 
 * @author takeshi
 */
public class TryCatchFinallyCompiler implements NodeCompiler {

    private final BSHTryStatement node;

    private final CompilationUnit unit;

    private Label startTry = new Label();

    private Label endTry = new Label();

    private Label handler = new Label(), outer = new Label();

    private boolean tryCatchVisited;

    private final Map<BSHFormalParameter, Tuple<BSHBlock, Label>> blocks = new LinkedHashMap<BSHFormalParameter, Tuple<BSHBlock, Label>>();

    private BSHBlock finallyBlock;

    /**
     * C'tor
     * 
     * @param node the node
     * @param unit the {@link CompilationUnit}
     */
    public TryCatchFinallyCompiler(BSHTryStatement node, CompilationUnit unit) {
        this.node = node;
        this.unit = unit;
    }

    /**
     * {@inheritDoc}
     */
    public void compile(MethodVisitor host, CompilationState state) throws CompilationFailedException {
        if (isTryBodyEmpty()) {
            finallyBlock = this.extractCatchBlocks();
            if (finallyBlock != null) {
                this.unit.compileNode(finallyBlock, host, state);
            }
            clearCompilerState();
            return;
        }
        Label after = new Label();
        int stack = state.getStackSize();
        host.visitLabel(startTry);
        visitTryCatchInstruction(host, state);

        // we must now compile the try { ... } block
        compileTryBody(host, state, after, stack);
        state.addLineNumber(host, this.node.getChild(0).getLineNumber());
        // now we must compile the exceptions
        for (Entry<BSHFormalParameter, Tuple<BSHBlock, Label>> entry : blocks.entrySet()) {
            this.compileCatchHandler(host, state, entry.getValue().tail(), after, entry.getKey(), entry.getValue().head());
        }
        state.assertStackSize(stack); // should not have changed the stack size
        if (finallyBlock != null) {
            host.visitLabel(outer);
            // must write the exception handler piece
            Scope scope = state.newScope(ScopeType.BLOCK);
            host.visitLabel(handler);
            scope.installStartLabel(host, this.unit.options());
            // the stack now contains a Throwable so we store on an unamed variable
            LocalVariable lv = scope.addLocalVariable(null, scope.getStartLabel(), JavaClassWrapper.wrapperFor(Throwable.class));
            // store the variable. Im not sure this is necessary, but javac does this. I could probably just leave it in the stack and use
            // it later
            host.visitVarInsn(Opcodes.ASTORE, lv.getOffset());
            // now we compile the finally bit
            this.unit.compileNode(finallyBlock, host, state);
            state.assertStackSize(stack); // should not have changed the stack size
            Scope exitScope = state.popScope();
            if (scope != exitScope) {
                throw new CompilationAssertionFailed("Poped scope != pushed scope");
            }
            exitScope.installEndLabel(host, this.unit.options());
            host.visitVarInsn(Opcodes.ALOAD, lv.getOffset());
            host.visitInsn(Opcodes.ATHROW);
        }
        host.visitLabel(after);
        clearCompilerState();
    }

    /**
     * Clears the compilation state. Should be called in case this try block is inside a finally clause - which gets compiled multiple times
     */
    public void clearCompilerState() {
        this.tryCatchVisited = false;
        this.endTry = new Label();
        this.startTry = new Label();
        handler = new Label();
        outer = new Label();
        blocks.clear();
    }

    /**
     * @return true if the try body is empty
     */
    private boolean isTryBodyEmpty() {
        return node.jjtGetNumChildren() == 0 || node.getChild(0).jjtGetNumChildren() == 0;
    }

    /**
     * Actually visits a try catch instruction
     * 
     * @param host {@link MethodVisitor}
     * @param state {@link CompilationState}
     * @throws CompilationFailedException e
     */
    public void visitTryCatchInstruction(MethodVisitor host, CompilationState state) throws CompilationFailedException {
        if (!tryCatchVisited) {
            visitInnerTryCatchInstructions(host, node.getChild(0), state);

            finallyBlock = extractCatchBlocks();

            for (Entry<BSHFormalParameter, Tuple<BSHBlock, Label>> entry : blocks.entrySet()) {
                TypeWrapper wrapper = this.parseExceptionType(entry.getKey(), state);
                host.visitTryCatchBlock(startTry, endTry, entry.getValue().tail(), wrapper.internalName());
            }

            if (finallyBlock != null) {
                // i must write an outer try { ... } catch (null null) { <finallyBlock> throw null }
                host.visitTryCatchBlock(startTry, outer, handler, null);
            }

            tryCatchVisited = true;
        }
    }

    /**
     * Visits try instructions that are inside the try body - the jvm requires inner tries to appear before the outer tries in the bytecode
     * 
     * @param host {@link MethodVisitor}
     * @param aNode {@link SimpleNode}
     * @param state {@link CompilationState}
     * @throws CompilationFailedException e
     */
    public void visitInnerTryCatchInstructions(MethodVisitor host, SimpleNode aNode, CompilationState state)
            throws CompilationFailedException {
        if (aNode instanceof BSHTryStatement) {
            TryCatchFinallyCompiler compiler = (TryCatchFinallyCompiler) unit.getNodeCompiler(aNode);
            compiler.visitTryCatchInstruction(host, state);
        } else {
            for (int i = 0; i < aNode.jjtGetNumChildren(); i++) {
                visitInnerTryCatchInstructions(host, aNode.getChild(i), state);
            }
        }
    }

    /**
     * Extracts all the catch blocks and stores in the {@link #blocks} field. Also returns the finally block
     * 
     * @return {@link BSHBlock}
     * @throws CompilationFailedException e
     */
    private BSHBlock extractCatchBlocks() throws CompilationFailedException {
        // first find the catch blocks and types that are caught
        for (int i = 1; i < node.jjtGetNumChildren(); i++) {
            SimpleNode maybeType = node.getChild(i);
            if (maybeType instanceof BSHFormalParameter) {
                final BSHFormalParameter formalParameter = (BSHFormalParameter) maybeType;
                // TypeWrapper exceptionType = parseExceptionType(formalParameter, state);
                BSHBlock body = (BSHBlock) node.getChild(++i);
                if (blocks.put(formalParameter, UnmodifiableTuple.newTuple(body, new Label())) != null) {
                    throw new CompilationFailedException(maybeType.getLineNumber(), "catching the same exception multiple times: "
                            + formalParameter.toString());
                }
            } else if (maybeType instanceof BSHBlock) {
                finallyBlock = (BSHBlock) maybeType;
            } else {
                throw new CompilationAssertionFailed("expected either BSHFormalParameter or BSHBlock but was: " + maybeType.getClass());
            }
        }
        return finallyBlock;
    }

    /**
     * Compiles the block that goes inside the catch block
     * 
     * @param host the {@link MethodVisitor}
     * @param state the current {@link CompilationState}
     * @param label the label for the start of the block
     * @param after the label to jump to after the catch block is done
     * @param key the variable declaration of the exception
     * @param value the catch block body
     * @throws CompilationFailedException e
     */
    private void compileCatchHandler(MethodVisitor host, CompilationState state, Label label, Label after, BSHFormalParameter key,
            BSHBlock value) throws CompilationFailedException {
        Scope newScope = state.newScope(ScopeType.BLOCK);
        newScope.installStartLabel(host, this.unit.options());
        host.visitLabel(label);
        TypeWrapper wrapper = this.parseExceptionType(key, state);
        LocalVariable variable = state.currentScope().addLocalVariable(key.name, label, wrapper);
        host.visitVarInsn(Opcodes.ASTORE, variable.getOffset());
        state.addLineNumber(host, key.getLineNumber());
        // now we compile the block properly
        int stack = state.getStackSize();
        unit.compileNode(value, host, state);
        alignStack(host, state, stack);
        Scope scope = state.popScope();
        scope.installEndLabel(host, this.unit.options());
        if (finallyBlock != null) {
            unit.compileNode(finallyBlock, host, state);
        }
        host.visitJumpInsn(Opcodes.GOTO, after);
    }

    /**
     * Compiles the try { ... } body
     * 
     * @param host the {@link MethodVisitor}
     * @param state the {@link CompilationState}
     * @param escape the label that breaks out of this try block
     * @param stack the stack size
     * @throws CompilationFailedException e
     */
    private void compileTryBody(MethodVisitor host, CompilationState state, Label escape, int stack)
            throws CompilationFailedException {
        Scope scope = state.newScope(ScopeType.BLOCK);
        scope.installStartLabel(host, this.unit.options());
        this.unit.compileNode(node.getChild(0), host, state);
        alignStack(host, state, stack);
        host.visitLabel(endTry);
        state.popScope();
        if (finallyBlock != null) {
            this.unit.compileNode(finallyBlock, host, state);
        }
        scope.installEndLabel(host, this.unit.options());
        host.visitJumpInsn(Opcodes.GOTO, escape);
    }

    /**
     * Aligns the stack size for two branching paths (the jvm requires this)
     * 
     * @param host the {@link MethodVisitor}
     * @param state the current state
     * @param stack the expected stack size after this method returns
     */
    private void alignStack(MethodVisitor host, CompilationState state, int stack) {
        if (state.getStackSize() != stack) {
            state.assertStackSize(stack + 1);
            TypeWrapper result = state.popStack();
            if (result.numberOfSlots() == 1) {
                host.visitInsn(Opcodes.POP);
            } else {
                host.visitInsn(Opcodes.POP2);
            }
        }
    }

    /**
     * Parses the exception types. Defaults to {@link Throwable}
     * 
     * @param maybeType it may or may not contain a type
     * @param state the {@link CompilationState}
     * @return a {@link TypeWrapper}
     * @throws CompilationFailedException e
     */
    private TypeWrapper parseExceptionType(BSHFormalParameter maybeType, CompilationState state) throws CompilationFailedException {
        if (maybeType.jjtGetNumChildren() == 1) {
            BSHType type = (BSHType) maybeType.getChild(0);
            return TypedVariableDeclarationCompiler.resolveTypeForNode(type, state.currentScope());
        } else {
            return JavaClassWrapper.wrapperFor(Throwable.class);
        }
    }

    /**
     * {@inheritDoc}
     */
    public TypeWrapper nodeResultType(Scope scope) throws CompilationFailedException {
        return JavaClassWrapper.VOID_WRAPPER;
    }

}
